# coding:utf-8
import os, torchvision
import torch.nn as nn
import numpy as np
import imageio
import matplotlib.pyplot as plt
from PIL import Image
import torch


def tensor2im(input_image, imtype=np.uint8, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]):
    assert len(input_image.shape) == 3
    if not isinstance(input_image, np.ndarray):
        if isinstance(input_image, torch.Tensor):
            image_tensor = input_image.data
        else:
            return input_image
        image_numpy = image_tensor.cpu().float().numpy()
        if image_numpy.shape[0] == 1:
            image_numpy = np.tile(image_numpy, (3, 1, 1))
        for i in range(len(mean)):
            image_numpy[i] = image_numpy[i] * std[i] + mean[i]
        image_numpy = image_numpy * 255
        image_numpy = np.transpose(image_numpy, (1, 2, 0))
    else:
        image_numpy = input_image
    return image_numpy.astype(imtype)


def save_img(im, path, size):
    """
    Parameters:
        im (tensor) --  a tuple of input tensor data
        path (str)  --  saving path
        size (int)  --  number of images in each row at maximum
    """
    im_grid = torchvision.utils.make_grid(im, size)
    im_numpy = tensor2im(im_grid)
    im_array = Image.fromarray(im_numpy)
    im_array.save(path)

